package tool.code.generater.util;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.DruidPooledConnection;
import com.alibaba.druid.util.StringUtils;
import tool.code.generater.config.JdbcConfig;

import java.sql.*;
import java.util.*;
import java.util.function.BiFunction;

public class DataSourceHelper {


    private final JdbcConfig jdbcConfig;

    public DataSourceHelper(JdbcConfig jdbcConfig) {
        this.jdbcConfig = jdbcConfig;
    }

    public DruidDataSource getDataSource() {
        DruidDataSource ds = new DruidDataSource();
        ds.setDriverClassName(jdbcConfig.getDriver());
        ds.setUrl(jdbcConfig.getJdbcUrl());
        ds.setUsername(jdbcConfig.getUsername());
        ds.setPassword(jdbcConfig.getPassword());
        return ds;
    }

    public List<Map<String, Object>> getColumns(String tableName) throws SQLException {
        return autoCloseWrapper((dpc, ps) -> {
            try {
                DatabaseMetaData metaData = dpc.getMetaData();
//                dpc.getCatalog() 数据库名称
//                dpc.getSchema() 逻辑库名称
                ResultSet columns = metaData.getColumns(dpc.getCatalog(), dpc.getSchema(), tableName, null);
                List<Map<String, Object>> result = new ArrayList<>();
                Map<String, Integer> colIndexMap = new HashMap<>();
                while (columns.next()) {
                    Map<String, Object> map = new LinkedHashMap<>();
                    String columnName = columns.getString("COLUMN_NAME");
                    map.put("name", columnName);
                    map.put("dataType", columns.getInt("DATA_TYPE"));
                    map.put("type", columns.getString("TYPE_NAME"));
                    map.put("size", columns.getInt("COLUMN_SIZE"));
                    map.put("remarks", columns.getString("REMARKS"));
                    colIndexMap.put(columnName, result.size());
                    result.add(map);
                }
                ResultSet primaryKeys = metaData.getPrimaryKeys(dpc.getCatalog(), dpc.getSchema(), tableName);
                while (primaryKeys.next()) {
                    String columnName = primaryKeys.getString("COLUMN_NAME");
                    String keySeq = String.valueOf(primaryKeys.getObject("KEY_SEQ"));
                    String pkName = primaryKeys.getString("PK_NAME");
                    if (colIndexMap.containsKey(columnName)) {
                        Map<String, Object> stringObjectMap = result.get(colIndexMap.get(columnName));
                        if (stringObjectMap != null) {
                            stringObjectMap.put("primaryKey", true);
                            stringObjectMap.put("keySeq", keySeq);
                            stringObjectMap.put("pkName", pkName);
                        }
                    }
                }
                return result;
            } catch (SQLException e) {
                throw SQLExceptionWrapper.wrap(e);
            }
        }, null);
    }

    public Map<String, String> getTablesComment() throws SQLException {
        return autoCloseWrapper((dpc, ps) -> {
            try {
                Map<String, String> map = new HashMap<>();
                DatabaseMetaData metaData = dpc.getMetaData();
                ResultSet rs = metaData.getTables(dpc.getCatalog(), dpc.getSchema(), null, null);
                while (rs.next()) {
                    String remarks = rs.getString("REMARKS");
                    String tableName = rs.getString("TABLE_NAME");
                    map.put(tableName, remarks);
                }
                return map;
            } catch (SQLException e) {
                throw SQLExceptionWrapper.wrap(e);
            }
        }, null);
    }

    public List<Map<String, Object>> doQuery(String sql) throws SQLException {
        return autoCloseWrapper((dpc, ps) -> {
            try {
                ResultSet resultSet = ps.executeQuery();
                List<Map<String, Object>> list = new ArrayList<Map<String, Object>>();
                ResultSetMetaData metaData = resultSet.getMetaData();
                int columnCount = metaData.getColumnCount();
                while (resultSet.next()) {
                    Map<String, Object> map = new LinkedHashMap<>();
                    for (int i = 0; i < columnCount; i++) {
                        map.put(metaData.getColumnLabel(i), resultSet.getObject(i + 1));
                    }
                    list.add(map);
                }
                return list;
            } catch (SQLException e) {
                throw SQLExceptionWrapper.wrap(e);
            }
        }, sql);
    }

    public int doUpdate(String sql) throws SQLException {
        if (StringUtils.isEmpty(sql)) {
            return 0;
        }
        return autoCloseWrapper((dpc, ps) -> {
            try {
                return ps.executeUpdate();
            } catch (SQLException e) {
                throw SQLExceptionWrapper.wrap(e);
            }
        }, sql);
    }

    public <T> T autoCloseWrapper(BiFunction<DruidPooledConnection, PreparedStatement, T> function, String sql) throws SQLException {
        DruidDataSource dataSource = null;
        DruidPooledConnection connection = null;
        PreparedStatement statement = null;
        try {
            dataSource = getDataSource();
            connection = dataSource.getConnection();
            if (sql != null && !sql.isEmpty()) {
                statement = connection.prepareStatement(sql);
            }
            return function.apply(connection, statement);
        } catch (SQLExceptionWrapper e) {
            throw e.unwrap();
        } finally {
            if (statement != null) {
                statement.close();
            }
            if (connection != null) {
                connection.close();
            }
            close(dataSource);
        }
    }

    public void close(DruidDataSource ds) {
        if (ds != null) {
            ds.close();
        }
    }

    public static class SQLExceptionWrapper extends RuntimeException {
        private static final long serialVersionUID = 1L;

        public SQLExceptionWrapper(String message) {
            super(message);
        }

        public SQLExceptionWrapper(String message, Throwable cause) {
            super(message, cause);
            assert cause instanceof SQLException;
        }

        public SQLExceptionWrapper(Throwable cause) {
            super(cause);
            assert cause instanceof SQLException;
        }

        public static SQLExceptionWrapper wrap(SQLException e) {
            return new SQLExceptionWrapper(e.getMessage(), e);
        }

        public SQLException unwrap() {
            return (SQLException) getCause();
        }
    }
}
